package main

import "fmt"

// 下沉操作，维护小顶堆的性质
func down(arr []int, i, size int) {
	for i < size/2 { // 只需要遍历到非叶子节点
		l := i*2 + 1                     // 左孩子的位置
		r := i*2 + 2                     // 右孩子的位置
		min := l                         // 假设左孩子是最小的
		if r < size && arr[r] < arr[l] { // 如果右孩子存在且比左孩子小，更新最小的位置
			min = r
		}
		if arr[i] > arr[min] { // 如果父节点大于最小的孩子，交换位置
			arr[i], arr[min] = arr[min], arr[i]
			i = min // 更新父节点的位置
		} else {
			break // 如果父节点小于等于最小的孩子，退出循环
		}
	}
}

// 堆排序方法，将数组从大到小排列
func HeapSort(arr []int) []int {
	n := len(arr)
	if n == 0 {
		return arr
	}
	for i := n/2 - 1; i >= 0; i-- { // 从最后一个非叶子节点开始，依次对每个节点进行下沉操作，直到根节点
		down(arr, i, n)
	}
	for i := n - 1; i >= 0; i-- { // 将堆中的元素依次取出放到数组中
		arr[0], arr[i] = arr[i], arr[0] // 将堆顶元素和数组末尾元素交换位置
		down(arr, 0, i)                 // 对堆顶元素进行下沉操作，维护剩余元素的堆性质
	}
	return arr // 返回排序后的数组
}

func main() {
	arr := []int{3, 5, 1, 4, 2} // 定义一个待排序的数组
	fmt.Println("The original array is:", arr)
	arr = HeapSort(arr) // 调用堆排序方法对数组进行排序
	fmt.Println("The sorted array is:", arr)

}
